modelopt/torch/export/plugins/vllm_fakequant_hf.py (1)

33-43: Docstring is outdated and references amax instead of quantizer state.

The docstring still describes extracting "amax values" but the implementation now saves the complete quantizer state dict and modelopt state. Update to reflect the actual behavior.

📝 Suggested docstring update

-    """Exports the torch model weights and amax values separately.
+    """Exports the torch model weights and quantizer state separately for vLLM fakequant.

     This function:
-    1. Extracts amax values for calibration
+    1. Extracts quantizer state dict and modelopt state
     2. Deletes all quantizer parameters from state dict to store only weights in original dtype
     3. Saves the model weights

     Args:
         model: The quantized model to export
-        export_dir: Directory to save the amax values
+        export_dir: Directory to save the model and quantizer state

     """

modelopt/torch/export/plugins/vllm_fakequant_megatron.py (1)

46-64: Comments reference "amax" but code now handles full quantizer state.

Several comments still reference "amax" (lines 46, 51, 63) but the code now handles the complete quantizer state dictionary. Consider updating for clarity.

📝 Suggested comment updates

-    # Gather all amax dicts to rank 0
+    # Gather all quantizer state dicts to rank 0
     world_size = torch.distributed.get_world_size()
     rank = torch.distributed.get_rank()

     if rank == 0:
-        # Rank 0 will collect all amax values
+        # Rank 0 will collect all quantizer state values
         all_quantizer_state_dicts = [None] * world_size
         torch.distributed.gather_object(quantizer_state_dict, all_quantizer_state_dicts, dst=0)
         ...
     else:
-        # Other ranks just send their amax values
+        # Other ranks send their quantizer state values
         torch.distributed.gather_object(quantizer_state_dict, None, dst=0)

examples/vllm_serve/vllm_reload_utils.py (1)

175-185: Docstring references non-existent parameter fuse_experts.

The docstring mentions fuse_experts parameter but the function only has state_dict and merge_mode parameters.

📝 Suggested docstring fix

     """
     Common implementation for converting quantizer state from HF to vLLM format.

     Args:
         state_dict: Input state dict
-        fuse_experts: Whether to fuse expert projections
         merge_mode: Mode to merge grouped values, "max_or_concat" or "require_identical"
+
+    Returns:
+        Converted state dict in vLLM format.
     """

examples/vllm_serve/fakequant_worker.py (2)

115-124: Consider adding a comment about weights_only=False security implications.

Using weights_only=False in torch.load is necessary for loading complex modelopt state, but it allows arbitrary code execution from untrusted files. The current code is fine for trusted checkpoints, but a brief comment noting this would be helpful for future maintainers.

💡 Optional: Add security note

         # Load on CPU to avoid failures when the checkpoint was saved from a different
         # GPU mapping
+        # Note: weights_only=False is required for modelopt state but should only be used
+        # with trusted checkpoint files.
         modelopt_state = torch.load(
             quant_config["modelopt_state_path"], weights_only=False, map_location="cpu"
         )

---

235-242: Asymmetric key validation is intentional but could use a comment.

The code raises an error when model keys are missing from the checkpoint but only warns when checkpoint has extra keys. This asymmetry makes sense (model requires all its quantizers to be loaded), but a brief comment explaining the rationale would improve clarity.

💡 Optional: Add clarifying comment

+            # Checkpoint may have extra keys (e.g., from PP sharding), but model must have
+            # all its quantizer keys present in the checkpoint for correct loading
             for key in checkpoint_quant_keys:
                 if key not in model_quant_keys:
                     print(f"Key {key} not found in model state dict, but exists in checkpoint")
             for key in model_quant_keys:
                 if key not in checkpoint_quant_keys:
                     raise ValueError(
                         f"Key {key} not found in checkpoint state dict, but exists in model"
                     )

examples/vllm_serve/hf_ptq_export.py (1)

41-41: Consider moving enable_huggingface_checkpointing() inside main().

Calling mto.enable_huggingface_checkpointing() at module level means it executes whenever this module is imported, which could cause unexpected side effects if the module is imported for other purposes (e.g., testing individual functions).

♻️ Proposed fix

-mto.enable_huggingface_checkpointing()
-
-
 def load_model(
     ...
+
+def main(args: argparse.Namespace):
+    if not torch.cuda.is_available():
+        raise OSError("GPU is required for inference.")
+
+    # Enable ModelOpt save/restore for HuggingFace models
+    mto.enable_huggingface_checkpointing()
+
+    random.seed(RAND_SEED)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/vllm_serve/hf_ptq_export.py` at line 41, The call to
mto.enable_huggingface_checkpointing() is currently at module import time and
should be moved inside main() to avoid side effects on import; locate the
module-level invocation of enable_huggingface_checkpointing(), remove it from
top-level, and add a call at the start of the main() function (before any model
loading or IO) so checkpointing is enabled only when main() runs.

modelopt/torch/quantization/plugins/custom.py (1)

117-117: Consider forwarding **kwargs to super().modelopt_post_restore().

The signature now accepts *args, **kwargs for compatibility, but the super() call at line 174 only passes prefix=prefix. If a caller passes model=..., it won't be forwarded to the base class QuantModule.modelopt_post_restore(), which now supports the model parameter for device discovery.

♻️ Proposed fix

         # If there are any other states, lets move them to the correct device
-        super().modelopt_post_restore(prefix=prefix)
+        super().modelopt_post_restore(prefix=prefix, **kwargs)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/quantization/plugins/custom.py` at line 117, The override
modelopt_post_restore(self, prefix: str = "", *args, **kwargs) must forward any
additional keyword args to the base implementation; update the call to
super().modelopt_post_restore(...) (the QuantModule.modelopt_post_restore path)
to pass prefix=prefix and also **kwargs (and *args if present) so parameters
such as model=... used for device discovery are propagated to the parent
implementation.

modelopt/torch/export/plugins/vllm_fakequant_megatron.py (1)

55-60: Variable shadowing in loop.

The loop variable quantizer_state_dict shadows the outer variable of the same name defined at line 42. This makes the code harder to read and could lead to confusion.

♻️ Proposed fix

         # Merge all quantizer state dicts into one
         merged_quantizer_state_dict = {}
-        for quantizer_state_dict in all_quantizer_state_dicts:
-            if quantizer_state_dict is not None:
-                merged_quantizer_state_dict.update(quantizer_state_dict)
+        for qstate_dict in all_quantizer_state_dicts:
+            if qstate_dict is not None:
+                merged_quantizer_state_dict.update(qstate_dict)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/export/plugins/vllm_fakequant_megatron.py` around lines 55 -
60, The loop is shadowing the outer variable quantizer_state_dict; rename the
loop variable (for example to iter_quantizer_state or q_state_dict) in the merge
block that iterates over all_quantizer_state_dicts and update any references in
that loop to use the new name so the outer quantizer_state_dict remains intact;
ensure merged_quantizer_state_dict.update(...) uses the renamed loop variable
and run tests to validate no other references rely on the old name.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

examples/vllm_serve/fakequant_worker.py (2)

115-117: ⚠️ Potential issue | 🔴 Critical

weights_only=False load still lacks the required safety justification.

This deserializes arbitrary pickled objects; add an inline safety justification (trusted/internal artifact) at the load site.

🔧 Proposed fix

+        # Safe: this file is an internally generated ModelOpt artifact from a trusted pipeline.
         modelopt_state = torch.load(
             quant_config["modelopt_state_path"], weights_only=False, map_location="cpu"
         )

As per coding guidelines, torch.load(..., weights_only=False) requires an inline comment justifying why it is safe.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/vllm_serve/fakequant_worker.py` around lines 115 - 117, The
torch.load call that sets modelopt_state using
quant_config["modelopt_state_path"] currently passes weights_only=False without
an inline safety justification; update the call site (the torch.load(...) that
assigns modelopt_state) to include a brief inline comment explaining why using
weights_only=False is safe (e.g., the file is a trusted/internal artifact
produced by our pipeline and contains no untrusted pickles), referencing the
trusted provenance and any validation performed.

---

223-223: ⚠️ Potential issue | 🟠 Major

Set weights_only explicitly when loading quantizer files.

Leaving this implicit makes behavior runtime-version dependent and can allow unsafe object deserialization on older PyTorch versions.

🔧 Proposed fix

-            saved_quant_dict = torch.load(quantizer_file_path, map_location="cpu")
+            saved_quant_dict = torch.load(
+                quantizer_file_path, map_location="cpu", weights_only=True
+            )

As per coding guidelines, unsafe torch.load usage must be explicitly constrained/justified rather than relying on implicit defaults.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/vllm_serve/fakequant_worker.py` at line 223, The torch.load call
that reads the quantizer file (saved_quant_dict =
torch.load(quantizer_file_path, map_location="cpu")) must explicitly pass
weights_only=True to avoid unsafe object deserialization and version-dependent
behavior; update the load invocation to include weights_only=True and, if
supporting older PyTorch versions, wrap the call in a try/except that falls back
to a safe alternative or raises a clear error indicating the required minimum
torch version—ensure references to quantizer_file_path and saved_quant_dict are
retained so the change is applied to the correct load site.

examples/vllm_serve/vllm_reload_utils.py (1)

200-204: ⚠️ Potential issue | 🟡 Minor

Validate merge_mode explicitly; don’t silently fallback.

A typo currently falls through to "max_or_concat" and can apply the wrong merge strategy without surfacing an error.

Proposed fix

-    merge_fn = (
-        _merge_values_require_identical
-        if merge_mode == "require_identical"
-        else _merge_values_by_max_or_concat
-    )
+    if merge_mode == "require_identical":
+        merge_fn = _merge_values_require_identical
+    elif merge_mode == "max_or_concat":
+        merge_fn = _merge_values_by_max_or_concat
+    else:
+        raise ValueError(f"Unsupported merge_mode: {merge_mode}")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/vllm_serve/vllm_reload_utils.py` around lines 200 - 204, The current
selection of merge_fn silently falls back to _merge_values_by_max_or_concat when
merge_mode is misspelled; change the logic that assigns merge_fn (the
conditional using merge_mode) to explicitly validate merge_mode and only accept
known values ("require_identical" and "max_or_concat"), selecting
_merge_values_require_identical or _merge_values_by_max_or_concat accordingly
and raising a clear ValueError (or similar) if merge_mode is invalid so typos
don’t silently change behavior.

examples/vllm_serve/vllm_reload_utils.py (1)

195-196: Remove stale fuse_experts docstring entry.

convert_dict_to_vllm has no fuse_experts parameter, so this is now misleading API documentation.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/vllm_serve/vllm_reload_utils.py` around lines 195 - 196, The
docstring for convert_dict_to_vllm incorrectly documents a non-existent
parameter "fuse_experts"; remove the stale "fuse_experts: Whether to fuse expert
projections" entry from the function/class docstring (search for
convert_dict_to_vllm in vllm_reload_utils.py) and update the remaining parameter
list so it only documents actual parameters (e.g., keep/verify merge_mode text
matches the function signature).

modelopt/torch/quantization/nn/modules/quant_module.py (1)

119-130: Consider a buffer fallback when parent has no parameters.

Line 120 only inspects named_parameters(). For buffer-only parent modules, device detection still fails even when a valid non-TensorQuantizer buffer exists. Adding a named_buffers() fallback would make restore more robust.

♻️ Proposed refactor

             # Look for any parameter in parent module (not just state_dict)
             for name, param in parent_module.named_parameters():
                 # Skip params that belong to TensorQuantizer submodules
                 param_parent_name = name.rsplit(".", 1)[0] if "." in name else ""
                 param_parent = (
                     parent_module.get_submodule(param_parent_name)
                     if param_parent_name
                     else parent_module
                 )
                 if not isinstance(param_parent, TensorQuantizer):
                     non_tq_param_or_buffer = param
                     break
+
+            # Fallback: some modules may be buffer-only
+            if non_tq_param_or_buffer is None:
+                for name, buf in parent_module.named_buffers():
+                    buf_parent_name = name.rsplit(".", 1)[0] if "." in name else ""
+                    buf_parent = (
+                        parent_module.get_submodule(buf_parent_name)
+                        if buf_parent_name
+                        else parent_module
+                    )
+                    if not isinstance(buf_parent, TensorQuantizer):
+                        non_tq_param_or_buffer = buf
+                        break

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/quantization/nn/modules/quant_module.py` around lines 119 -
130, The parent-module device detection loop in quant_module.py currently only
checks parent_module.named_parameters() and thus misses modules that expose only
buffers; update the logic in the restore/device-detection area (the loop that
iterates parent_module.named_parameters() and sets non_tq_param_or_buffer) to
also fall back to parent_module.named_buffers() when no suitable parameter is
found, applying the same TensorQuantizer parent check (using param_parent_name
and parent_module.get_submodule(...)) so a non-TensorQuantizer buffer can be
selected; ensure the code still prefers parameters over buffers but uses a
buffer if no parameter exists.

examples/vllm_serve/vllm_reload_utils.py (1)

140-163: Consider adding defensive check for empty key_value_pairs.

If key_value_pairs is empty, values[0] will raise an IndexError. While the caller (convert_dict_to_vllm) currently only calls this when len(key_value_pairs) > 1, adding a guard would make the function more robust against future misuse.

🛡️ Proposed defensive check

 def _merge_values_by_max_or_concat(merged_key: str, key_value_pairs: list[tuple[str, Any]]) -> Any:
     """
     Merge values by taking max for amax, concatenating for others.
     Used for quantizer state weights (tensor values).
     """
+    if not key_value_pairs:
+        raise ValueError(f"Cannot merge empty key_value_pairs for '{merged_key}'")
     values = [value for _, value in key_value_pairs]

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/vllm_serve/vllm_reload_utils.py` around lines 140 - 163, The
function _merge_values_by_max_or_concat should defensively handle an empty
key_value_pairs input to avoid IndexError on values[0]; add an early guard at
the top of the function that checks if not key_value_pairs and raise a clear
ValueError (or return an appropriate sentinel) with a message referencing
merged_key; keep the rest of logic unchanged—this makes the helper robust even
if convert_dict_to_vllm or other callers misuse it.